{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some examples of different types of ERA5 data that are publicly available for download.\n",
    "f1 = \"era5/output_selected_vars.nc\"\n",
    "f2 = \"era5/2d29cff3c6faebc67ea7ee432b0cf937.nc\"\n",
    "f3 = \"era5/data_stream-oper_stepType-instant.nc\"\n",
    "f4 = \"era5/data_stream-oper_stepType-accum.nc\"\n",
    "\n",
    "d1 = xr.open_dataset(f1, engine='netcdf4')\n",
    "d2 = xr.open_dataset(f2, engine='netcdf4')\n",
    "d3 = xr.open_dataset(f3, engine='netcdf4')\n",
    "d4 = xr.open_dataset(f4, engine='netcdf4')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def standardize_spatial(ds):\n",
    "    ds = ds.rename({'latitude': 'lat', 'longitude': 'lon'})\n",
    "    if ds.lat[0] > ds.lat[-1]:\n",
    "        ds = ds.reindex(lat=ds.lat[::-1])\n",
    "    return ds\n",
    "\n",
    "d2 = standardize_spatial(d2)\n",
    "d3 = standardize_spatial(d3)\n",
    "d4 = standardize_spatial(d4)\n",
    "\n",
    "target_times = pd.to_datetime([\n",
    "    \"20xx-xx-xxT18:00:00\",\n",
    "    \"20xx-xx-xxT00:00:00\",\n",
    "    \"20xx-xx-xxT06:00:00\"\n",
    "])\n",
    "\n",
    "d2 = d2.sel(valid_time=target_times)\n",
    "d3 = d3.sel(valid_time=target_times)\n",
    "d4 = d4.sel(valid_time=target_times)\n",
    "\n",
    "ds_final = xr.Dataset()\n",
    "\n",
    "ds_final = ds_final.expand_dims(batch=[0])\n",
    "\n",
    "ds_final.coords['lon'] = d1['lon']\n",
    "ds_final.coords['lat'] = d1['lat']\n",
    "ds_final.coords['level'] = d2['pressure_level']\n",
    "ds_final.coords['time'] = ('time', target_times)\n",
    "\n",
    "ds_final.coords['datetime'] = (('batch', 'time'), np.array([target_times]))\n",
    "\n",
    "\n",
    "ds_final['geopotential_at_surface'] = d1['geopotential_at_surface']\n",
    "ds_final['land_sea_mask'] = d1['land_sea_mask']\n",
    "\n",
    "var_mapping_3d = {\n",
    "    'z': 'geopotential',\n",
    "    't': 'temperature',\n",
    "    'u': 'u_component_of_wind',\n",
    "    'v': 'v_component_of_wind',\n",
    "    'w': 'vertical_velocity',\n",
    "    'q': 'specific_humidity'\n",
    "}\n",
    "\n",
    "for src_name, dst_name in var_mapping_3d.items():\n",
    "    da = d2[src_name]\n",
    "    da = da.rename({'valid_time': 'time', 'pressure_level': 'level'})\n",
    "    da = da.expand_dims('batch', axis=0)\n",
    "    ds_final[dst_name] = da.transpose('batch', 'time', 'level', 'lat', 'lon')\n",
    "\n",
    "d3 = d3.rename({'valid_time': 'time'})\n",
    "d3 = d3.expand_dims('batch').assign_coords(batch=[0])\n",
    "\n",
    "ds_final['2m_temperature'] = d3['t2m'].transpose('batch', 'time', 'lat', 'lon')\n",
    "ds_final['mean_sea_level_pressure'] = d3['msl'].transpose('batch', 'time', 'lat', 'lon')\n",
    "ds_final['10m_u_component_of_wind'] = d3['u10'].transpose('batch', 'time', 'lat', 'lon')\n",
    "ds_final['10m_v_component_of_wind'] = d3['v10'].transpose('batch', 'time', 'lat', 'lon')\n",
    "\n",
    "d4 = d4.rename({'valid_time': 'time'})\n",
    "d4 = d4.expand_dims('batch').assign_coords(batch=[0])\n",
    "\n",
    "ds_final['total_precipitation_6hr'] = d4['tp'].transpose('batch', 'time', 'lat', 'lon')\n",
    "ds_final['toa_incident_solar_radiation'] = d4['tisr'].transpose('batch', 'time', 'lat', 'lon')\n",
    "\n",
    "ds_final.attrs['description'] = 'Unified ERA5 dataset with standardized dimensions'\n",
    "ds_final.attrs['created_at'] = pd.Timestamp.now().isoformat()\n",
    "\n",
    "print(ds_final)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_final.to_netcdf(\"...\", engine='netcdf4')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
